import math

import numpy as np

# 这个文件是和算法对接的 转换成指令用的
# 前面说的函数就是这里用的
from agent.agent import Agent
from env.env_cmd import EnvCmd, Point
from env.env_def import RED_AIRPORT_ID, UnitStatus, UnitType


class RedAgent(Agent):
    def __init__(self, name, all_args, config, **kwargs):
        super().__init__(name, config["side"])
        self.num_agents = all_args.num_agents
        self.num_enemys = all_args.num_enemys

    def reset(self):
        pass

    # env_runner.py _agent_step
    def step(self, sim_time, action_n, obs, **kwargs):
        cmd_list = []
        obs_red = obs["red"]
        is_red = False
        if obs_red["airports"][0]["AIR"] > 0 and sim_time > 0:
            # 机场控制飞机起飞
            return [
                EnvCmd.make_takeoff_areapatrol(
                    RED_AIRPORT_ID,
                    min(1, obs_red["airports"][0]["AIR"]),
                    UnitType.A2A,
                    100000,
                    0,
                    8000,
                    0,
                    20000,
                    20000,
                    1000,
                    7200,
                )
            ], is_red

        if sim_time >= 200:
            is_red = True
            # obs是观测的
            input_obs = np.zeros(
                (self.num_agents * 2 + self.num_enemys * 2), dtype=np.float32
            )
            self_a2a = []
            enemy_a2a = []
            i = 0
            my_a2a = 0
            # 遍历红方的units，看是否是飞机
            for unit in obs["red"]["units"]:
                if unit["LX"] == UnitType.A2A:
                    input_obs[i] = unit["X"] / 100000
                    i += 1
                    input_obs[i] = unit["Y"] / 100000
                    i += 1
                    # input_obs[0][i] = unit["WP"].get("170", 0) / 10
                    # i += 1
                    my_a2a += 1
                    self_a2a.append(unit)
            i = self.num_agents * 2
            en_a2a = 0
            for unit in obs["blue"]["units"]:
                if unit["LX"] == UnitType.A2A:
                    input_obs[i] = unit["X"] / 100000
                    i += 1
                    input_obs[i] = unit["Y"] / 100000
                    i += 1
                    enemy_a2a.append(unit)
                    en_a2a += 1
            actions = np.argmax(action_n, axis=1)
            # print("选择动作：", actions)
            trap_list = [[0.2, 0], [0, 0.2], [-0.2, 0], [0, -0.2]]
            target_x, target_y = (
                input_obs[self.num_agents * 2],
                input_obs[self.num_agents * 2 + 1],
            )
            possible_points = [
                (target_x + trap[0], target_y + trap[1]) for trap in trap_list
            ]
            for i in range(my_a2a):
                i_unit = self_a2a[i]
                # direction = ((actions[i]) / len(action_n[0])) * 2 * np.pi
                # 计算角度
                # 提取坐标
                x1, y1 = input_obs[i * 2], input_obs[i * 2 + 1]

                # 计算每个点到 (x1, y1) 的距离
                distances = [
                    np.linalg.norm([x1 - px, y1 - py]) for px, py in possible_points
                ]

                # 选择距离最近的点
                min_index = distances.index(min(distances))
                x2, y2 = possible_points[min_index]
                x1, y1 = input_obs[i * 2], input_obs[i * 2 + 1]
                x2, y2 = (
                    input_obs[self.num_agents * 2] + trap_list[actions[i]][0],
                    input_obs[self.num_agents * 2 + 1] + trap_list[actions[i]][1],
                )

                # 计算向量
                vector = np.array([x2 - x1, y2 - y1])
                distance = np.linalg.norm(vector)
                # 计算角度（弧度）
                angle_rad = np.arctan2(vector[1], vector[0])
                direction = angle_rad
                # print(
                #     "第{i}个飞机位置为:{input_obs[2*i]},{input_obs[2*i+1]},目标位置为:{input_obs[-2]},{input_obs[-1]}，方向为:",
                #     np.cos(direction),
                #     " ",
                #     np.sin(direction),
                # )
                # print("x, y: ", np.cos(direction), np.sin(direction))
                cmd_list.append(
                    EnvCmd.make_areapatrol(
                        i_unit["ID"],
                        i_unit["X"]
                        + min(distance * 100000, 5000) * np.cos(direction).item(),
                        i_unit["Y"]
                        + min(distance * 100000, 5000) * np.sin(direction).item(),
                        i_unit["Z"],
                        0,
                        1000,
                        1000,
                        min(max(100, distance * 100000 / 40), 1000),
                        7200,
                    )
                )

        return cmd_list, is_red
